from tensorflow.keras.datasets import mnist
from tensorflow.keras.utils import to_categorical

from tensorflow.keras import layers
from tensorflow.keras.models import Model,load_model
from tensorflow.keras import optimizers
from tensorflow.keras.losses import categorical_crossentropy






(x_train,y_train),(x_test,y_test) = mnist.load_data()
X_train = x_train.reshape((60000,28,28,1))/255.
X_test = x_test.reshape((10000,28,28,1))/255.

y_train = to_categorical(y_train,num_classes=10)
y_test = to_categorical(y_test,num_classes=10)


tempetature = 3
inputs = layers.Input((28,28,1))
x = layers.Conv2D(16,3,activation='relu')(inputs)
x = layers.Conv2D(16,3,strides=2,activation='relu')(x)
x = layers.Conv2D(32,5,activation='relu')(x)
x = layers.Conv2D(32,5,activation='relu')(x)
x = layers.Flatten()(x)
x = layers.Dense(60,activation='relu')(x)
x = layers.Dense(10,activation='softmax')(x)
x = layers.Lambda(lambda t:t/tempetature)(x)
student_less = Model(inputs,x)
teacher = load_model('teacher_model/')
student_less.compile(optimizer=optimizers.SGD(momentum=0.9,nesterov=True),loss=categorical_crossentropy,metrics=['accuracy'])
student_less.fit(X_train,teacher.predict(X_train)/tempetature,batch_size=128,epochs=1,verbose=2)
student_less.evaluate(X_test,y_test/tempetature)
student_less.save("student_less/")